from typing import Sequence
import numbers

import random
from PIL import Image, ImageOps

import numpy as np


def hflip(img):
    return img.transpose(Image.FLIP_LEFT_RIGHT)


def vflip(img):
    return img.transpose(Image.FLIP_TOP_BOTTOM)


def crop(img, top, left, height, width):
    return img.crop((left, top, left + width, top + height))


def rotate(img, angle, resample=False, expand=False, center=None, fill=0):

    if isinstance(fill, int):
        fill = tuple([fill] * 3)

    return img.rotate(angle, resample, expand, center, fillcolor=fill)


def pad(img, padding, fill=0, padding_mode='constant'):
    r"""Pad the given PIL Image on all sides with specified padding mode and fill value.

    Args:
        img (PIL Image): Image to be padded.
        padding (int or tuple): Padding on each border. If a single int is provided this
            is used to pad all borders. If tuple of length 2 is provided this is the padding
            on left/right and top/bottom respectively. If a tuple of length 4 is provided
            this is the padding for the left, top, right and bottom borders
            respectively.
        fill: Pixel fill value for constant fill. Default is 0. If a tuple of
            length 3, it is used to fill R, G, B channels respectively.
            This value is only used when the padding_mode is constant
        padding_mode: Type of padding. Should be: constant, edge, reflect or symmetric. Default is constant.

            - constant: pads with a constant value, this value is specified with fill

            - edge: pads with the last value on the edge of the image

            - reflect: pads with reflection of image (without repeating the last value on the edge)

                       padding [1, 2, 3, 4] with 2 elements on both sides in reflect mode
                       will result in [3, 2, 1, 2, 3, 4, 3, 2]

            - symmetric: pads with reflection of image (repeating the last value on the edge)

                         padding [1, 2, 3, 4] with 2 elements on both sides in symmetric mode
                         will result in [2, 1, 1, 2, 3, 4, 4, 3]

    Returns:
        PIL Image: Padded image.
    """

    if not isinstance(padding, (numbers.Number, tuple)):
        raise TypeError('Got inappropriate padding arg')
    if not isinstance(fill, (numbers.Number, str, tuple)):
        raise TypeError('Got inappropriate fill arg')
    if not isinstance(padding_mode, str):
        raise TypeError('Got inappropriate padding_mode arg')

    if isinstance(padding, Sequence) and len(padding) not in [2, 4]:
        raise ValueError("Padding must be an int or a 2, or 4 element tuple, not a " +
                         "{} element tuple".format(len(padding)))

    assert padding_mode in ['constant', 'edge', 'reflect', 'symmetric'], \
        'Padding mode should be either constant, edge, reflect or symmetric'

    if padding_mode == 'constant':
        if img.mode == 'P':
            palette = img.getpalette()
            image = ImageOps.expand(img, border=padding, fill=fill)
            image.putpalette(palette)
            return image

        return ImageOps.expand(img, border=padding, fill=fill)
    else:
        if isinstance(padding, int):
            pad_left = pad_right = pad_top = pad_bottom = padding
        if isinstance(padding, Sequence) and len(padding) == 2:
            pad_left = pad_right = padding[0]
            pad_top = pad_bottom = padding[1]
        if isinstance(padding, Sequence) and len(padding) == 4:
            pad_left = padding[0]
            pad_top = padding[1]
            pad_right = padding[2]
            pad_bottom = padding[3]

        if img.mode == 'P':
            palette = img.getpalette()
            img = np.asarray(img)
            img = np.pad(img, ((pad_top, pad_bottom), (pad_left, pad_right)), padding_mode)
            img = Image.fromarray(img)
            img.putpalette(palette)
            return img

        img = np.asarray(img)
        # RGB image
        if len(img.shape) == 3:
            img = np.pad(img, ((pad_top, pad_bottom), (pad_left, pad_right), (0, 0)), padding_mode)
        # Grayscale image
        if len(img.shape) == 2:
            img = np.pad(img, ((pad_top, pad_bottom), (pad_left, pad_right)), padding_mode)

        return Image.fromarray(img)


class RandomHorizontalFlip(object):

    def __init__(self, prob=0.5):
        self.prob = prob

    def __call__(self, lr, hr):
        if random.random() < self.prob:
            return hflip(lr), hflip(hr)
        return lr, hr


class RandomVerticalFlip(object):

    def __init__(self, prob=0.5):
        self.prob = prob

    def __call__(self, lr, hr):
        if random.random() < self.prob:
            return vflip(lr), vflip(hr)
        return lr, hr


class RandomRot90(object):

    def __init__(self, prob=0.5):
        self.prob = prob

    def __call__(self, lr, hr):
        if random.random() < self.prob:
            return rotate(lr, 90), rotate(hr, 90)
        return lr, hr


class RandomCrop(object):

    def __init__(self, patch_size, scale):
        self.patch_size = patch_size
        self.scale = scale

    def __call__(self, lr, hr):

        iw, ih = lr.size

        tp = self.patch_size
        ip = tp // self.scale

        ix = random.randrange(0, iw - ip + 1)
        iy = random.randrange(0, ih - ip + 1)
        tx, ty = self.scale * ix, self.scale * iy
        lr = crop(lr, iy, ix, ip, ip)
        hr = crop(hr, ty, tx, tp, tp)
        return lr, hr

    def __repr__(self):
        return self.__class__.__name__ + '(patch_size={0}, scale={1})'.format(self.patch_size, self.scale)


class ToTensor:

    def __call__(self, lr, hr):
        lr = np.array(lr).transpose((2, 0, 1)).astype(np.float32)
        hr = np.array(hr).transpose((2, 0, 1)).astype(np.float32)
        return lr, hr